# 这是一个开头
# 人员：sulei
# 开发时间：18/11/2020下午8:09
# 文件名：2D_VAEs_train_main.py
# 开发工具：PyCharm


import time
import torch.optim as optim
# from torch.optim import lr_scheduler
import matplotlib.pyplot as plt
from collections import OrderedDict
from tqdm import tqdm
import  torch
import numpy as np
from torch.utils.data import DataLoader
from torch import nn
from functions import *
from model import VAE


def train(model, train_loader, optimizer, epoch, quiet, grad_clip=None):
    model.train()

    if not quiet:
        pbar = tqdm(total=len(train_loader.dataset))
    losses = OrderedDict()
    # t= 0
    for x in tqdm(train_loader,total=38):
        if x.dtype is not torch.float32:
            x = x.detach().float()
        x = x.cuda()
        out,sigma = model.loss(x)
        optimizer.zero_grad()  # 手动梯度置零
        out['loss'].backward()  # 反向传播

        if grad_clip:
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
        optimizer.step()

        desc = "Epoch{}".format(epoch)

        for k, v in out.items():  # k, v  =('a','b')  a  recon_loss,kl_loss,loss
            if k not in losses:
                losses[k] = []
            losses[k].append(v.item())
            avg_loss = np.mean(losses[k][-50:])
            desc += ",{} {:.4f}".format(k,avg_loss)
        desc += ",{} {:.4f}".format('sigma',sigma)
        if not quiet:
            pbar.set_description(desc)  # 与34行   ,进度条
            pbar.update(x.shape[0])
    if not quiet:
        pbar.close()
    return losses

def eval_loss(model, data_loader, quiet):
    model.eval()
    total_losses = OrderedDict()
    with torch.no_grad():
        for x in data_loader:
            if x.dtype is not torch.float32:
                x = x.detach().float()
            x = x.cuda()
            out,sigma = model.loss(x)
            for k, v in out.items():
                total_losses[k] = total_losses.get(k, 0) + v.item() * x.shape[0]

        desc = 'Test '
        for k in total_losses.keys():
            total_losses[k] /= len(data_loader.dataset)
            desc += ",{} {:.4f}".format(k,total_losses[k])

    return total_losses

def train_epochs(model, train_loader, test_loader, train_args, quiet=False):

    epochs, lr = train_args['epochs'], train_args['lr']
    grad_clip = train_args.get('grad_clip', None)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    # loss_func = torch.nn.MSELoss()

    #学习率自调整
    # exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=learning_rate_decay_every, gamma=.8)


    train_losses, test_losses = OrderedDict(), OrderedDict()

    print("开始训练：")
    for epoch in range(epochs):
        model.train()
        train_loss = train(model, train_loader, optimizer, epoch, quiet, grad_clip)

        test_loss = eval_loss(model, test_loader, quiet)

        # exp_lr_scheduler.step()

        for k in train_loss.keys():
            if k not in train_losses:
                train_losses[k] = []
                test_losses[k] = []
            train_losses[k].extend(train_loss[k])
            # test_losses[k].append(test_loss[k])
            # 改动后，10*test_loss
            test_losses[k].append(test_loss[k])
    return train_losses, test_losses

if __name__ == "__main__":

    t_start = time.time()

    load_path ='./data/skl/datasets_4800X600_time/datasets_4800x600_timeorder.npy'
    # load_path ='./data/datasets_4800X150.npy'


    train_datasets,test_datasets = data_loader(load_path)

    train_loader = DataLoader(train_datasets,batch_size=64,shuffle=True)

    test_loader = DataLoader(test_datasets,batch_size=128,shuffle=False)

    model = VAE().cuda()

    train_losses, test_losses = train_epochs(model, train_loader, test_loader,
                                             dict(epochs=100, lr=0.0001), quiet=False)

    train_losses = np.stack((train_losses['loss'], train_losses['recon_loss'], train_losses['kl_loss']), axis=1)

    test_losses = np.stack((test_losses['loss'], test_losses['recon_loss'], test_losses['kl_loss']), axis=1)

    print(f'Final -ELBO: {test_losses[-1, 0]:.4f}, Recon Loss: {test_losses[-1, 1]:.4f}, '
          f'KL Loss: {test_losses[-1, 2]:.4f}')

    plot_vae_training_plot(train_losses, test_losses, title='vae_notime_traing',
                           fname= 'vae_notime_traing')

    train_datasets = np.load(load_path)
    train_datasets = torch.from_numpy(train_datasets)
    train_datasets = train_datasets.detach().float().cuda()
    z_datasets = model.sample(train_datasets)
    z_datasets = z_datasets.cpu().detach().numpy()
    #
    save_path ='./data/out/datasets_4800X400_time.npy'

    np.save(save_path,z_datasets)


    t_end = time.time()

    print(t_end - t_start)







